import sys, fitz
import shutil,os
from cnocr import CnOcr
from poprogress import simple_progress
 
def pyMuPDF_fitz(pdfPath, imagePath):#pdf转成图片
    pdfDoc = fitz.open(pdfPath)
    page = pdfDoc[0]#只转第一页
    rotate = int(0)
    zoom_x = 1.3
    zoom_y = 1.3
    mat = fitz.Matrix(zoom_x, zoom_y).prerotate(rotate)
    pix = page.get_pixmap(matrix=mat, alpha=False)          
    pix.save ('pdf.png')   # 保存图片

def pdfRename(pdfPath,pdfName,res,stuNames,stuNames_dict,stuCard_dict):#重命名  
    name=""#记录识别到的学生姓名
    pdfNewName = ""#记录目标文件名
    record=False#当遇到姓名或者身份证号时，将此变量设置为true，表示下一个内容要记录
    exitFile=False#记录完身份证后，将此变量设置为true，表示获取完成，可以退出当前文件的处理
    for i in range(len(res)):#遍历图片识别结果，是一个由字典组成的列表
        for k,v in res[i].items():#遍历识别出的每一项数据，数据是一个字典
            if k=="text":
                if record:
                    if exitFile:
                        cardid=str(v).strip()#去除识别出来的身份证前后的空格
                        if len(cardid)==0:#如果身份证为空，识别下一个
                            continue
                        pdfNewName=cardid+name
                        if pdfNewName in stuNames:#如果文件名在name.txt文件中，直接修改
                            os.rename(pdfPath+pdfName,pdfPath+pdfNewName+".pdf")
                        elif str(v) in stuNames_dict.keys():#如果文件名不在，但是身份证号码正确，则修改姓名，以身份证号码+修改后的姓名作为文件名重命名
                            os.rename(pdfPath+pdfName,pdfPath+str(v)+stuNames_dict[str(v)]+".pdf")
                            shutil.copy(pdfPath+str(v)+stuNames_dict[str(v)]+".pdf",pdfPath+"待检查\姓名识别失败")
                        elif name in stuCard_dict:
                            os.rename(pdfPath+pdfName,pdfPath+stuCard_dict[name]+name+".pdf")
                            shutil.copy(pdfPath+stuCard_dict[name]+name+".pdf",pdfPath+"待检查\身份证识别失败")
                        else:#如果身份证号码和姓名都不正确，不修改，提示手动修改
                            print("{}识别失败，识别为{},自动修改失败，请手动修改。".format(pdfName,pdfNewName))
                        return
                    name=str(v)#记录姓名
                    record=False
                elif v=="姓名":
                    record=True
                elif v=="身份证号":
                    record=True
                    exitFile=True
                else:
                    record=False

if __name__ == "__main__":
    pdfPath="d:/pdf/"#设置存放pdf文件的文件夹路径,路径最后的"/"不要漏
    if os.path.exists(pdfPath+"待检查"):
        shutil.rmtree(pdfPath+"待检查")#删除待检查文件夹
    filenames=os.listdir(pdfPath)#遍历存放pdf的文件夹，获取所有pdf文件名
    os.makedirs(pdfPath+"待检查")
    os.makedirs(pdfPath+"待检查\姓名识别失败")
    os.makedirs(pdfPath+"待检查\身份证识别失败")
    stuNames=[]
    stuNames_dict={}
    stuCard_dict={} 
    with open(os.path.join(sys.path[0], 'name.txt'), "r", encoding='utf-8') as f:  #打开文本
        for n in f.readlines():
            stuNames.append(n.strip('\n'))      #去除文本中的换行符
    for n in stuNames:#建立身份信息字典，key为身份证，value为姓名，用于校对和纠正识别的结果
        stuNames_dict[n[0:18]]=n[18:]
    for k,v in stuNames_dict.items():#建立身份信息字典，key为姓名，value为身份证，用于校对和纠正识别的结果
        if v not in stuCard_dict.keys():#如果姓名在字典中不存在，则加入字典，如果已经存在，则表示有学生姓名重复，此时不能根据姓名确定身份证号码，予以删除
            stuCard_dict[v]=k
        else:
            del stuCard_dict[v]
    for pdfName in simple_progress(filenames,desc="识别进度："):
        if os.path.isfile("pdf.png"):
            os.remove("pdf.png")
        pyMuPDF_fitz(pdfPath+pdfName, "./")#调用函数，将pdf的第一页转成图片
        ocr = CnOcr()
        res = ocr.ocr("pdf.png")#识别图片
        pdfRename(pdfPath,pdfName,res,stuNames,stuNames_dict,stuCard_dict)  #重命名
    print("转换结束。\n身份证或姓名识别错误的文档复制到了对应文件夹下，请检查。如有问题，在pdf目录下手动修改相应文件，识别正确可忽略。")